﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;

public enum TCP_TABLE_TYPE : int
{
    TCP_TABLE_BASIC_LISTENER,
    TCP_TABLE_BASIC_CONNECTIONS,
    TCP_TABLE_BASIC_ALL,
    TCP_TABLE_OWNER_PID_LISTENER,
    TCP_TABLE_OWNER_PID_CONNECTIONS,
    TCP_TABLE_OWNER_PID_ALL,
    TCP_TABLE_OWNER_MODULE_LISTENER,
    TCP_TABLE_OWNER_MODULE_CONNECTIONS,
    TCP_TABLE_OWNER_MODULE_ALL
}

public enum TCP_CONNECTION_STATE : int
{
    CLOSED = 1,
    LISTENING,
    SYN_SENT,
    SYN_RCVD,
    ESTABLISHED,
    FIN_WAIT_1,
    FIN_WAIT_2,
    CLOSE_WAIT,
    CLOSING,
    LAST_ACK,
    TIME_WAIT,
    DELETE_TCP
}

public enum UDP_TABLE_TYPE : int
{
    UDP_TABLE_BASIC,
    UDP_TABLE_OWNER_PID,
    UDP_TABLE_OWNER_MODULE
}

namespace Kernel.Connection
{
    public class GTcpConnectionTableHelper   //定义TCP的基本数据类型供之后调用
    {
        [DllImport("Ws2_32.dll")]
        static extern ushort ntohs(ushort netshort);

        [DllImport("iphlpapi.dll", SetLastError = true)]
        static extern uint GetExtendedTcpTable(IntPtr pTcpTable, ref int dwOutBufLen, bool sort, int ipVersion, TCP_TABLE_TYPE tblClass, int reserved);

        [StructLayout(LayoutKind.Sequential)]
        public struct MIB_TCPROW_OWNER_PID
        {
            public uint state;
            public uint localAddr;
            public byte localPort1;
            public byte localPort2;
            public byte localPort3;
            public byte localPort4;
            public uint remoteAddr;
            public byte remotePort1;
            public byte remotePort2;
            public byte remotePort3;
            public byte remotePort4;
            public int owningPid;

            public ushort LocalPort
            {
                get
                {
                    return BitConverter.ToUInt16(new byte[2] { localPort2, localPort1 }, 0);
                }
            }

            public ushort RemotePort
            {
                get
                {
                    return BitConverter.ToUInt16(new byte[2] { remotePort2, remotePort1 }, 0);
                }
            }
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct MIB_TCPTABLE_OWNER_PID
        {
            public uint dwNumEntries;
            MIB_TCPROW_OWNER_PID table;
        }

        public static string GetIpAddress(long ipAddrs)
        {
            try
            {
                System.Net.IPAddress ipAddress = new System.Net.IPAddress(ipAddrs);
                return ipAddress.ToString();
            }
            catch { return ipAddrs.ToString(); }

        }

        public static ushort GetTcpPort(int tcpPort)
        {
            return ntohs((ushort)tcpPort);
        }

        public static MIB_TCPROW_OWNER_PID[] GetAllTcpConnections()
        {
            MIB_TCPROW_OWNER_PID[] tcpConnectionRows;
            int AF_INET = 2;    // IPv4
            int buffSize = 0;

            // use WinAPI GetExtendedTcpTable to query all active tcp connection information
            uint ret = GetExtendedTcpTable(IntPtr.Zero, ref buffSize, true, AF_INET, TCP_TABLE_TYPE.TCP_TABLE_OWNER_PID_ALL, 0);
            if (ret != 0 && ret != 122) // 122 means insufficient buffer size
            {
                throw new Exception("Error occurred when trying to query tcp table, return code: " + ret);
            }
            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);

            try
            {
                ret = GetExtendedTcpTable(buffTable, ref buffSize, true, AF_INET, TCP_TABLE_TYPE.TCP_TABLE_OWNER_PID_ALL, 0);
                if (ret != 0)
                {
                    throw new Exception("Error occurred when trying to query tcp table, return code: " + ret);
                }

                // get the number of entries in the table
                MIB_TCPTABLE_OWNER_PID table = (MIB_TCPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_TCPTABLE_OWNER_PID));

                IntPtr rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(table.dwNumEntries));
                tcpConnectionRows = new MIB_TCPROW_OWNER_PID[table.dwNumEntries];

                for (int i = 0; i < table.dwNumEntries; i++)
                {
                    MIB_TCPROW_OWNER_PID tcpRow = (MIB_TCPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_TCPROW_OWNER_PID));
                    tcpConnectionRows[i] = tcpRow;
                    rowPtr = (IntPtr)((long)rowPtr + Marshal.SizeOf(tcpRow));
                }
            }
            finally
            {
                // free memory
                Marshal.FreeHGlobal(buffTable);
            }
            return tcpConnectionRows;
        }
    }

    public class GUdpConnectionTableHelper   //定义UDP的基本数据类型供之后调用
    {
        [DllImport("Ws2_32.dll")]
        static extern ushort ntohs(ushort netshort);

        [DllImport("iphlpapi.dll", SetLastError = true)]
        static extern uint GetExtendedUdpTable(IntPtr pTcpTable, ref int dwOutBufLen, bool sort, int ipVersion, UDP_TABLE_TYPE tblClass, int reserved);

        [StructLayout(LayoutKind.Sequential)]
        public struct MIB_UDPROW_OWNER_PID
        {
            public uint state
            {
                get { return 15; }
            }
            public uint localAddr;
            public byte localPort1;
            public byte localPort2;
            public byte localPort3;
            public byte localPort4;
            //public uint remoteAddr;
            //public byte remotePort1;
            //public byte remotePort2;
            //public byte remotePort3;
            //public byte remotePort4;
            public int owningPid;

            public ushort LocalPort
            {
                get
                {
                    return BitConverter.ToUInt16(new byte[2] { localPort2, localPort1 }, 0);
                }
            }

            public ushort RemotePort
            {
                get
                {
                    return 0;
                }
            }
            public uint remoteAddr
            {
                get { return 0; }
           }
        }

        [StructLayout(LayoutKind.Sequential)]
        public struct MIB_UDPTABLE_OWNER_PID
        {
            public uint dwNumEntries;
            MIB_UDPROW_OWNER_PID table;
        }

        public static string GetIpAddress(long ipAddrs)
        {
            try
            {
                System.Net.IPAddress ipAddress = new System.Net.IPAddress(ipAddrs);
                return ipAddress.ToString();
            }
            catch { return ipAddrs.ToString(); }

        }

        public static ushort GetUdpPort(int udpPort)
        {
            return ntohs((ushort)udpPort);
        }

        public static MIB_UDPROW_OWNER_PID[] GetAllUdpConnections()
        {
            MIB_UDPROW_OWNER_PID[] udpConnectionRows;
            int AF_INET = 2;    // IPv4
            int buffSize = 0;

            // use WinAPI GetExtendedTcpTable to query all active tcp connection information
            uint ret = GetExtendedUdpTable(IntPtr.Zero, ref buffSize, true, AF_INET, UDP_TABLE_TYPE.UDP_TABLE_OWNER_PID, 0);

            if (ret != 0 && ret != 122) // 122 means insufficient buffer size
            {
                throw new Exception("Error occurred when trying to query tcp table, return code: " + ret);
            }
            IntPtr buffTable = Marshal.AllocHGlobal(buffSize);

            try
            {
                ret = GetExtendedUdpTable(buffTable, ref buffSize, true, AF_INET, UDP_TABLE_TYPE.UDP_TABLE_OWNER_PID, 0);
                if (ret != 0)
                {
                    throw new Exception("Error occurred when trying to query tcp table, return code: " + ret);
                }

                // get the number of entries in the table
                MIB_UDPTABLE_OWNER_PID table = (MIB_UDPTABLE_OWNER_PID)Marshal.PtrToStructure(buffTable, typeof(MIB_UDPTABLE_OWNER_PID));

                IntPtr rowPtr = (IntPtr)((long)buffTable + Marshal.SizeOf(table.dwNumEntries));
                udpConnectionRows = new MIB_UDPROW_OWNER_PID[table.dwNumEntries];

                for (int i = 0; i < table.dwNumEntries; i++)
                {
                    MIB_UDPROW_OWNER_PID udpRow = (MIB_UDPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_UDPROW_OWNER_PID));
                    udpConnectionRows[i] = udpRow;
                    rowPtr = (IntPtr)((long)rowPtr + Marshal.SizeOf(udpRow));
                }
            }
            finally
            {
                // free memory
                Marshal.FreeHGlobal(buffTable);
            }
            return udpConnectionRows;
        }
    }

    public class GetPIDsConnectionInfo //获取进程TCP UDP连接信息
    {
        public struct ConnectionInfo
        {
            public Process ThisProcess
            { get; set; }
            public uint _state;
            public string State
            {
                get
                {
                    switch(_state)
                    {
                        case 1: return "关闭";
                        case 2: return "监听";
                        case 3: return "异步发送";
                        case 4: return "异步接收";
                        case 5: return "建立";
                        case 6: return "FIN等待1";
                        case 7: return "FIN等待2";
                        case 8: return "等待关闭";
                        case 9: return "关闭中";
                        case 10: return "请求连接";
                        case 11: return "TIME_WAIT";
                        case 12: return "DELETE_TCP";
                        default: return "无";
                    }
                }
            }
            public string LocalAddr { get; set; }
            public string RemoteAddr { get; set; }
            public int PID { get; set; }
            public ushort LocalPort { get;set; }
            public ushort RemotePort {get;set;}
            public string PIDName { get; set; }
            public string LinkType { get; set; }
        }
        public List<ConnectionInfo> ConnectionInfoList;
        private GTcpConnectionTableHelper.MIB_TCPROW_OWNER_PID[] _TcpProgressInfoTable;
        public GTcpConnectionTableHelper.MIB_TCPROW_OWNER_PID[] TcpProgressInfoTable
        {
            get { return _TcpProgressInfoTable; }
            private set { _TcpProgressInfoTable = value; }
        }
        private GUdpConnectionTableHelper.MIB_UDPROW_OWNER_PID[] _UdpProgressInfoTable;
        public GUdpConnectionTableHelper.MIB_UDPROW_OWNER_PID[] UdpProgressInfoTable
        {
            get { return _UdpProgressInfoTable; }
            private set { _UdpProgressInfoTable = value; }
        }
        public int ConnectionPIDNumber { get; set; }
        private Process[] ProcessList=Process.GetProcesses();

        public void GetInfo()
        {
            ConnectionPIDNumber = 0;
            ConnectionInfoList = new List<ConnectionInfo>(); 
            TcpProgressInfoTable = GTcpConnectionTableHelper.GetAllTcpConnections();
            foreach (var tcp in TcpProgressInfoTable)
            {
                ConnectionInfo adds = new ConnectionInfo
                {
                    PID = tcp.owningPid,
                    _state = tcp.state,
                    LocalAddr = GTcpConnectionTableHelper.GetIpAddress(tcp.localAddr),
                    LocalPort = tcp.LocalPort,
                    RemoteAddr = GTcpConnectionTableHelper.GetIpAddress(tcp.remoteAddr),
                    RemotePort = tcp.RemotePort,
                    LinkType="TCP"
                };
                try
                {
                    adds.ThisProcess = Process.GetProcessById(adds.PID);
                    adds.PIDName = adds.ThisProcess.ProcessName;
                }
                catch { adds.ThisProcess = null; adds.PIDName = "未知"; }
                ConnectionInfoList.Add(adds);
            }
            UdpProgressInfoTable = GUdpConnectionTableHelper.GetAllUdpConnections();
            foreach (var udp in UdpProgressInfoTable)
            {
                ConnectionInfo adds = new ConnectionInfo
                {
                    PID = udp.owningPid,
                    _state = udp.state,
                    LocalAddr = GTcpConnectionTableHelper.GetIpAddress(udp.localAddr),
                    LocalPort = udp.LocalPort,
                    RemoteAddr = GTcpConnectionTableHelper.GetIpAddress(udp.remoteAddr),
                    RemotePort = udp.RemotePort,
                    LinkType = "UDP"
                };
                try
                {
                    adds.ThisProcess = Process.GetProcessById(adds.PID);
                    adds.PIDName = adds.ThisProcess.ProcessName;
                }
                catch { adds.ThisProcess = null; adds.PIDName = "未知"; }
                ConnectionInfoList.Add(adds);
            }
            List<string> namelist=new List<string>();
            foreach(var con in ConnectionInfoList)
            {
                if (!namelist.Contains(con.PIDName))
                {
                    namelist.Add(con.PIDName);
                    ConnectionPIDNumber++;
                }
            }
        }
    }
}